"""
该文件仅用于通过LSTM对训练模型进行测试
"""
import os
import time
import config
import numpy as np
import tensorflow as tf
import LSTM.GRU as GRU
from LSTM import data_manage as dm


os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'


SEQ_MAX_LEN = config.SEQ_MAX_LEN    # 序列最大长度
N_HIDDEN_UNITS = config.N_HIDDEN_UNITS  # 隐藏节点数目
BATCH_SIZE = config.LSTM_BATCH_SIZE  # 批处理大小
CLASS_NUMBERS = config.CLASS_NUMBERS    # 分类数目
FILE_DATA_TEST = config.FILE_DATA_TEST  # 分割后的临时测试数据集文件
FILE_SAVE = config.FILE_SAVE    # 存储运行结果目录
MODEL_PATH = FILE_SAVE + '/Model/model.ckpt-400'    # 模型保存路径


CLASS_NAMES = {
    0: 'Grasp', 1: 'Turn Left', 2: 'Turn Right', 3: 'Slide Down', 4: 'Slide Left', 5: 'Slide Right',
    6: 'Wave', 7: 'Flip'
}


# 标签数值化
def one_hot(label):
    label_num = len(label)
    new_label = label.reshape(label_num)
    n_values = np.max(new_label) + 1   # 看结果取出最大的
    return np.eye(n_values)[np.array(new_label, dtype=np.int32)]


if __name__ == '__main__':

    # 设置网络输入的占入符
    with tf.name_scope('input'):
        X = tf.placeholder("float", [None, SEQ_MAX_LEN, 5], name='x_input')
        Y = tf.placeholder("float", [None, CLASS_NUMBERS], name='y_input')
        true_length = tf.placeholder(tf.int32, [None], name='length_input')

    # 设置权重和偏置参数
    weights = {
        'in': tf.Variable(tf.truncated_normal([5, N_HIDDEN_UNITS]), name='w_in'),
        'out': tf.Variable(tf.truncated_normal([N_HIDDEN_UNITS, CLASS_NUMBERS]), name='w_out')
    }
    biases = {
        'in': tf.Variable(tf.constant(0.1, shape=[N_HIDDEN_UNITS, ]), name='b_in'),
        'out': tf.Variable(tf.constant(0.1, shape=[CLASS_NUMBERS, ]), name='b_out')
    }

    # 读取测试数据集
    test_dic = dm.read_data_txt(FILE_DATA_TEST)
    test_x = test_dic['x']
    test_y = test_dic['y']
    test_len = test_dic['length']
    test_num = test_dic['number']
    test_y_hot = one_hot(test_y)

    logits = GRU.GRU(X, weights, biases, true_length)
    prediction = tf.nn.softmax(logits)
    pre_num = tf.argmax(prediction, 1, name="output")
    correct_pred = tf.equal(pre_num, tf.argmax(Y, 1))  # shape(128,1)
    accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))
    saver = tf.train.Saver()    # 加载模型持久化类
    time_start = time.time()

    with tf.Session() as sess:
        saver.restore(sess, MODEL_PATH)     # 加载ckpt模型
        output = sess.run(pre_num, feed_dict={
            X: np.array(test_x)[0: BATCH_SIZE],
            Y: np.array(test_y_hot)[0: BATCH_SIZE],
            true_length: np.array(test_len)[0: BATCH_SIZE]
        })
        correct = [0] * CLASS_NUMBERS
        number = [0] * CLASS_NUMBERS
        for i in range(BATCH_SIZE):
            jud = False
            if test_y[i] == output[i]:
                jud = True
                correct[int(test_y[i])] += 1
            number[int(test_y[i])] += 1
            print(test_y[i], output[i], jud)
        correct_overall = 0
        for i in range(len(correct)):
            print('%d\t ACC: %.4f' % (i, correct[i] / number[i]))
            correct_overall += correct[i]
        print('ACC OVERALL: %.4f' % (correct_overall / BATCH_SIZE))

        # 程序运行计时
        time_end = time.time()
        input_time = time_end - time_start
        print('Totally time: %.4f' % input_time)
